﻿from openai import OpenAI
import base64
import os
from typing import List, Dict
from datetime import datetime

# -------------------------- 1. 核心参数配置 --------------------------
API_KEY = "sk-PjfrbJ5yZ1kAeSH6875c4fD5551f4f4b8e839d28Da91A6Db"  # 替换为你的正确密钥
MODEL_ID = "gemini-2.5-flash-image-preview"
# 待生成的图片描述列表（每条描述将与最终地址关联）

OUTPUT_DIR = "./gemini_image_output"
GEN_PARAMS = {
    "temperature": 0.7,
    "modalities": ["text", "image"]
}

# -------------------------- 2. 工具函数 --------------------------
def init_output_dir() -> None:
    """初始化保存目录"""
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)
        print(f"✅ 保存目录创建成功：{os.path.abspath(OUTPUT_DIR)}")
    else:
        print(f"ℹ️ 保存目录已存在：{os.path.abspath(OUTPUT_DIR)}")

def base64_to_image(base64_str: str, save_path: str) -> bool:
    """Base64解码为图片并保存"""
    try:
        image_data = base64.b64decode(base64_str.strip())
        with open(save_path, "wb") as f:
            f.write(image_data)
        return True
    except Exception as e:
        print(f"❌ 图片保存失败：{str(e)}")
        return False

# -------------------------- 3. 单图生成函数（返回绝对路径） --------------------------
def generate_single_image(
    client: OpenAI,
    prompt: str,
    save_filename: str
) -> str or None:
    """成功返回图片绝对路径，失败返回None"""
    full_prompt = f"Generate a 16:9 HD image strictly according to the description: {prompt}"
    try:
        print(f"\n📌 处理描述：{prompt[:50]}...")
        # 调用API
        response = client.chat.completions.create(
            model=MODEL_ID,
            messages=[{"role": "user", "content": full_prompt}],
            **GEN_PARAMS
        )

        # 解析图片数据
        message = response.choices[0].message
        if not hasattr(message, "multi_mod_content") or not message.multi_mod_content:
            print(f"❌ 无multi_mod_content字段（描述：{prompt[:30]}...）")
            return None

        # 查找有效Base64数据
        image_base64 = None
        for idx, mod_item in enumerate(message.multi_mod_content):
            if "inline_data" in mod_item and mod_item["inline_data"]:
                inline_data = mod_item["inline_data"]
                if "data" in inline_data and inline_data["data"].strip():
                    image_base64 = inline_data["data"].strip()
                    print(f"✅ 第{idx+1}个元素找到图片数据（长度：{len(image_base64)}）")
                    break

        if not image_base64:
            print(f"❌ 未找到有效图片数据（描述：{prompt[:30]}...）")
            return None

        # 保存并返回绝对路径
        absolute_path = os.path.abspath(save_filename)
        if base64_to_image(image_base64, absolute_path):
            print(f"✅ 生成成功：{absolute_path}")
            return absolute_path
        return None

    except Exception as e:
        print(f"❌ 生成异常（描述：{prompt[:30]}...）：{str(e)}")
        return None

# -------------------------- 4. 批量生成函数（返回“描述-地址”关联列表） --------------------------
def batch_generate_images(IMAGE_PROMPTS) -> List[Dict[str, str]]:
    """
    返回格式：
    [
        {"prompt": "图片描述1", "image_path": "图片1绝对路径"},
        {"prompt": "图片描述2", "image_path": "图片2绝对路径"}
    ]
    仅包含生成成功的记录
    """
    # 前置检查：OpenAI SDK版本
    try:
        import openai
        from packaging.version import Version
        if Version(openai.__version__) < Version("1.0.0"):
            print(f"⚠️ 请升级SDK：pip install -U openai（当前：{openai.__version__}）")
            return []
    except ImportError:
        print(f"⚠️ 未安装SDK：pip install -U openai")
        return []

    # 初始化客户端和目录
    init_output_dir()
    try:
        client = OpenAI(api_key=API_KEY, base_url="https://aihubmix.com/v1")
        print(f"✅ 客户端初始化成功（模型：{MODEL_ID}）")
    except Exception as e:
        print(f"❌ 客户端失败：{str(e)}")
        return []

    # 核心：存储“描述-地址”关联的列表
    generated_result: List[Dict[str, str]] = []
    total = len(IMAGE_PROMPTS)
    success_count = 0

    print(f"\n🚀 批量任务启动（共{total}张图片）")
    for idx, prompt in enumerate(IMAGE_PROMPTS, 1):
        # 生成唯一文件名
        safe_keyword = prompt.replace(" ", "_").replace(",", "").replace("：", "").replace("/", "")[:20]
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        save_path = os.path.join(OUTPUT_DIR, f"image_{safe_keyword}.png")

        # 生成图片并关联描述
        image_absolute_path = generate_single_image(client, prompt, save_path)
        if image_absolute_path:
            # 新增“描述-地址”字典到列表
            generated_result.append({
                "prompt": prompt,       # 原始描述
                "image_path": image_absolute_path  # 绝对路径
            })
            success_count += 1

    # 打印总结（展示关联关系）
    print(f"\n" + "="*60)
    print(f"📊 批量任务完成！")
    print(f"总计：{total}条描述 | 成功：{success_count}条 | 失败：{total-success_count}条")
    print(f"\n📋 成功生成的“描述-地址”列表：")
    for i, item in enumerate(generated_result, 1):
        print(f"\n{i}. 描述：{item['prompt']}")
        print(f"   地址：{item['image_path']}")
    print("="*60)

    return generated_result  # 返回关联列表

# -------------------------- 5. 主执行入口（示例如何使用关联列表） --------------------------
if __name__ == "__main__":
    # 执行批量生成，获取“描述-地址”关联列表
    image_prompts: List[str] = [
        "一只柯基犬戴着圣诞帽，趴在堆满礼物的地毯上，背景是暖光圣诞树，治愈系插画",
        "卡通风格太空猫，穿宇航员服漂浮在火星表面，身边有旋转小行星"
    ]
    prompt_path_list = batch_generate_images(image_prompts)

    # 示例1：遍历列表，根据描述筛选图片（如筛选“柯基”相关图片）
    if prompt_path_list:
        print(f"\n💡 示例1：筛选含“柯基”的图片")
        corgi_images = [item for item in prompt_path_list if "柯基" in item["prompt"]]
        if corgi_images:
            for item in corgi_images:
                print(f" - 描述：{item['prompt']} | 地址：{item['image_path']}")
        else:
            print(f" - 未找到含“柯基”的图片")

    # 示例2：遍历列表，批量处理图片（如批量去水印、上传等）
    if prompt_path_list:
        print(f"\n💡 示例2：批量处理所有成功图片")
        for idx, item in enumerate(prompt_path_list, 1):
            print(f"\n处理第{idx}张：")
            print(f"  描述：{item['prompt']}")
            print(f"  原地址：{item['image_path']}")
            # 这里可添加实际处理逻辑（如去水印）
            # no_watermark_path = item['image_path'].replace(".png", "_no_watermark.png")
            # remove_watermark(item['image_path'], no_watermark_path)
            # print(f"  去水印后地址：{no_watermark_path}")
    else:
        print(f"\n💡 无成功生成的图片，无需后续处理")